import os

import pandas as pd

from app_config import get_engine_ts
from app_config import get_pro


def calculate_volatility(ts_code, df1, df_q):
    list1 = df_q['cal_date'].tolist()
    import pandas as pd

    # 假设df1是你的dataframe，list1是你的集合列表
    # 假设df1中存在一列名为'trade_date'

    # 1. 获取dataframe的最后一行索引
    last_row_index = df1.shape[0] - 1

    # 2. 获取dataframe的'trade_date'列的所有元素列表
    trade_dates = df1['trade_date'].tolist()

    # 3. 找出不在trade_date列的元素
    new_dates = list(set(list1) - set(trade_dates))

    # 4. 创建一个新的DataFrame，包含这些新的元素
    new_rows = pd.DataFrame([df1.iloc[last_row_index]] * len(new_dates), columns=df1.columns)
    new_rows['trade_date'] = new_dates

    # 5. 将这些数据作为新的行插入到dataframe中
    df1 = pd.concat([df1, new_rows], ignore_index=True)

    # 6. 确保所有其他列用末行的数据填充
    for column in df1.columns:
        if column != 'trade_date':
            df1.loc[df1.index[-len(new_dates):], column] = df1.iloc[last_row_index][column]
    #
    # # 现在，df1已经被更新，new_dates中的元素被追加到了trade_date列中，其他列使用末行的数据填充
    # if ts_code == '600941.SH':
    #     df1.to_excel("zfile/沪深300_2024季度4_预期/" + ts_code + ".xlsx")
    return df1['total_mv'].mean()


def calc_hs300(strStartDate='', strEndDate='', strSeasonDate='', str_filePre='', cur_000300Date='', pre_000300Date=''):
    if not os.path.exists(str_filePre):
        os.makedirs(str_filePre)
        print(f"文件夹 '{str_filePre}' 创建成功")

    engine = get_engine_ts()

    query_hfd = f"""
           SELECT ts_code, symbol, name, market, industry, list_date FROM `stock_basic`
           """
    db_stock_basic = pd.read_sql_query(query_hfd, engine)
    ''' 
        指数样本空间由同时满足以下条件的非 ST、*ST 沪深 A 股和红
        筹企业发行的存托凭证组成：
         科创板证券、创业板证券：上市时间超过一年。
         其他证券：上市时间超过一个季度，除非该证券自上市以
        来日均总市值排在前 30 位。 
    '''
    stock_basic_st = db_stock_basic[~db_stock_basic['name'].str.contains('ST')]
    stock_basic = stock_basic_st[(
                                         (stock_basic_st['market'].isin(['科创板', '创业板']))
                                         &
                                         (stock_basic_st['list_date'] <= strStartDate)
                                 ) | (
                                         (stock_basic_st['market'] == '主板')
                                         &
                                         (stock_basic_st['list_date'] <= strSeasonDate)
                                 )]
    stock_basic.to_excel(str_filePre + '01stock_base.xlsx')

    query = f"""
       SELECT ts_code,trade_date, amount FROM `daily`
       WHERE trade_date >= '{strStartDate}'
       """

    # 执行查询并将结果转换为DataFrame
    db_daily = pd.read_sql_query(query, engine)
    db_daily['amount'] = db_daily['amount'] / 1000
    grouped_df = db_daily.groupby('ts_code')['amount'].mean().reset_index()
    grouped_df.columns = ['ts_code', '日均成交额']
    # 按日均成交额 列降序排序
    df_sort_amount = grouped_df.sort_values(by='日均成交额', ascending=False).reset_index(drop=True)
    list_extend_ts_code = df_sort_amount.head(30)['ts_code'].tolist()

    list_index_ts_code = stock_basic['ts_code'].tolist()
    list_index_ts_code.extend(list_extend_ts_code)

    df_stock_basic = db_stock_basic[db_stock_basic['ts_code'].isin(list_index_ts_code)]

    # 构建查询语句
    query = f"""
        SELECT ts_code,trade_date, total_mv FROM `daily_basic`
        WHERE trade_date >= '{strStartDate}'
        """
    # 执行查询并将结果转换为DataFrame
    db_daily = pd.read_sql_query(query, engine)
    db_daily['total_mv'] = db_daily['total_mv'] / 10000

    df_q = get_pro().trade_cal(exchange='SSE', is_open='1',
                               start_date=strStartDate,
                               end_date=strEndDate,
                               fields='cal_date')

    '''计算年化波动率 好像没什么用'''
    # 测试函数（需根据实际情景选择相应的DataFrame）
    # 假设我们已经加载了数据在df中，并且它是按ts_code分组的
    volatility_dict = {}
    # 遍历计算波动率
    for stock_code, group_df in db_daily.groupby('ts_code'):
        volatility_dict[stock_code] = calculate_volatility(stock_code, group_df, df_q)
    # 将字典转换为 DataFrame
    grouped_df_total_mv = pd.DataFrame(list(volatility_dict.items()), columns=['ts_code', '日均总市值'])

    merge = pd.merge(df_stock_basic, grouped_df, on='ts_code', how='left')
    pd_merge = pd.merge(merge, grouped_df_total_mv, on='ts_code', how='left')

    pre_000300_index = get_pro().index_weight(index_code='000300.SH', start_date=pre_000300Date,
                                              end_date=pre_000300Date)
    # pre_000300_index.rename(columns={'con_code': 'ts_code'}, inplace=True)
    my_set = pre_000300_index['con_code'].tolist()

    # 按日均成交额 列降序排序
    df_sorted = pd_merge.sort_values(by='日均成交额', ascending=False).reset_index(drop=True)
    # 获取前 90% 的数据
    num_rows = len(df_sorted)
    # df_sorted['i_avgAmount'] = df_sorted.index.astype(str) + '/' + str(int(num_rows * 0.5))
    df_sorted['i_日均成交'] = df_sorted.index.astype(int) + 1
    df_sorted.loc[df_sorted['ts_code'].isin(my_set), 'i_日均成交'] = 0
    df_sorted = df_sorted.sort_values(by='i_日均成交').reset_index(drop=True)
    top_50_percent = df_sorted.head(int(num_rows * 0.5))

    top_50_percent = top_50_percent.sort_values('日均总市值', ascending=False).reset_index(drop=True)
    top_50_percent['i_日均市值'] = top_50_percent.index
    df_result_hs300 = top_50_percent.head(400)
    # print("开始合并")
    df_000300 = get_pro().index_weight(index_code='000300.SH', start_date=cur_000300Date, end_date=cur_000300Date)
    df_000300.rename(columns={'con_code': 'ts_code'}, inplace=True)
    df_000300.rename(columns={'trade_date': 'list_date'}, inplace=True)
    df_000300['symbol'] = df_000300['ts_code'].str.split('.').str[0]
    df_000300.drop(columns=['weight'], inplace=True)

    df_000300 = pd.merge(df_000300, df_sorted[['ts_code', 'i_日均成交', '日均成交额']], on='ts_code', how='left')
    df_000300.to_excel(str_filePre + cur_000300Date + '-000300.xlsx')

    # 获取 dataframe1 中的所有列
    columns_dataframe1 = df_result_hs300.columns

    # 创建缺失列并填充为 '*'
    for column in columns_dataframe1:
        if column not in df_000300.columns:
            df_000300[column] = '**'

    # 将 dataframe2 的列顺序调整为与 dataframe1 一致
    df_000300 = df_000300[columns_dataframe1]

    # 追加 dataframe2 到 dataframe1
    result = pd.concat([df_result_hs300, df_000300], ignore_index=True)

    result.to_excel(str_filePre + '沪深300_.xlsx')


if __name__ == '__main__':
    # # """2024年四季度"""
    # strStartDate = '20231101'
    # strSeasonDate = '20240731'
    # strEndDate = '20241031'
    # str_filePre = 'zfile/沪深300_2024季度4_预期_20241126/'
    # pre_000300 = '20240628'  # 上一期沪深300
    # str_000688Date = '20240628'
    # calc_hs300(strStartDate=strStartDate, strEndDate=strEndDate, strSeasonDate=strSeasonDate, str_filePre=str_filePre,
    #            cur_000300Date=str_000688Date,
    #            pre_000300Date=pre_000300)

    # # """2025二季度"""
    # strStartDate = '20240501'
    # strSeasonDate = '20250131'
    # strEndDate = '20250429'
    # str_filePre = 'zfile/hs300_2025_02season/'
    # pre_date = '20250228'  # 上一期沪深300
    # cur_date = '20250228'
    # calc_hs300(strStartDate=strStartDate, strEndDate=strEndDate, strSeasonDate=strSeasonDate, str_filePre=str_filePre,
    #            cur_000300Date=cur_date,
    #            pre_000300Date=pre_date)

    # """2025二季度"""
    strStartDate = '20241101'
    strEndDate = '20251031'
    strSeasonDate = '20250731'
    str_filePre = 'zfile/hs300_2025_04season/'
    pre_date = '20250630'  # 上一期沪深300
    cur_date = '20250630'
    calc_hs300(strStartDate=strStartDate,
               strEndDate=strEndDate,
               strSeasonDate=strSeasonDate,
               str_filePre=str_filePre,
               cur_000300Date=cur_date,
               pre_000300Date=pre_date)
